"""
Author: Morphlng
Date: 2024-04-08 21:31:59
LastEditTime: 2025-03-18 12:54:51
LastEditors: Morphlng
Description: Driving Gym environment
FilePath: /DrivingGym/src/driving_gym/environment/env.py
"""

from __future__ import annotations

import logging
from copy import deepcopy

import gymnasium as gym

from driving_gym.data.data_provider import DataProvider
from driving_gym.environment.agent.rl_agent import RLAgent
from driving_gym.environment.scenario.scenario_manager import ScenarioManager
from driving_gym.misc.viz import Visualizer
from driving_gym.simulation.adapter_factory import AdapterFactory

logger = logging.getLogger(__name__)


class DrivingGym(gym.Env):
    def __init__(self, config: dict):
        self.config = deepcopy(config)

        # Adapter
        self.adapter_type = self.config.get("adapter_type", None)
        self.adapter_config = self.config["adapter"]
        self.adapter = AdapterFactory.create(self.adapter_config, self.adapter_type)

        # Scenario
        self.scenario_config = self.config["scenarios"]
        self.scenario_manager = ScenarioManager(self.scenario_config, self.adapter)
        self.data_provider = DataProvider(self.adapter)

        # Agent
        self.agent_config = self.config["agents"]
        for name, agent_config in self.agent_config.items():
            agent_config["actor_id"] = name

        self.agents = {
            name: RLAgent(agent_config, self.adapter)
            for name, agent_config in self.agent_config.items()
        }
        self.observation_space = gym.spaces.Dict(
            {name: agent.observation_space for name, agent in self.agents.items()}
        )
        self.action_space = gym.spaces.Dict(
            {name: agent.action_space for name, agent in self.agents.items()}
        )

        # Episode information
        self._initialized = False
        self.scenario = None
        self.num_episode = 0
        self.num_step = 0

        # Render
        self.visualizer = Visualizer(self.config.get("visualizer", {}))
        self.visualizer.parse_config(self.agent_config)

    def reset(self, seed=None, options=None):
        if not self._initialized:
            self.adapter.initialize()
            self._initialized = True

        self.scenario = self.scenario_manager.get_scenario()
        self.scenario.reset()
        for name, agent in self.agents.items():
            agent.reset(self.scenario)

        self.data_provider.reset()
        self.num_episode += 1
        self.num_step = 0

        # TODO: tick?
        for _ in range(10):
            self.data_provider.collect_frame_data(self.adapter.tick())

        obs, info = {}, {}
        for name, agent in self.agents.items():
            nframe = agent.obs_handler.frame_stack
            nstep = agent.obs_handler.frame_skip + 1
            snapshots = self.data_provider.get_frames(nframe, nstep)
            obs[name] = agent.get_obs(snapshots)

        return obs, info

    def step(self, action_dict: dict):
        self.num_step += 1

        # TODO: STEP_TICKS and ACTION_TICKS
        for name, action in action_dict.items():
            control = self.agents[name].act(action, self.num_step)
            self.adapter.apply_control(name, control)
        self.scenario.run_step()
        self.data_provider.collect_frame_data(self.adapter.tick())

        obs, reward, terminate, info = {}, {}, {}, {}
        for name in action_dict:
            agent = self.agents[name]
            nframe = agent.obs_handler.frame_stack
            nstep = agent.obs_handler.frame_skip + 1
            snapshots = self.data_provider.get_frames(nframe, nstep)
            obs[name] = agent.get_obs(snapshots)
            reward[name], reward_detail = agent.get_reward(self.data_provider.last_frame)
            terminate[name], cause = agent.get_done(self.data_provider.last_frame)
            info[name] = reward_detail
            if terminate[name]:
                logger.info(f"Agent {name} is terminated due to {cause}")

        terminate["__all__"] = all(terminate.values())
        if self.num_step >= self.scenario.max_steps:
            truncate = {name: True for name in action_dict}
            logger.info("Episode is truncated due to max steps")
        else:
            truncate = {name: False for name in action_dict}
        truncate["__all__"] = all(truncate.values())

        return obs, reward, terminate, truncate, info

    def render(self):
        if self.visualizer.total_sensors_per_agent > 0:
            self.visualizer.render(self.data_provider.last_frame)

    def close(self):
        self.adapter.cleanup()
        self.visualizer.close()
        self._initialized = False
        self.scenario = None
        self.num_episode = 0
        self.num_step = 0
